import os
from typing import Callable
import math
import numpy as np
import torch
from dora import DoraStatus
import time
IMAGE_WIDTH = 1920
IMAGE_HEIGHT = 1080
DEVICE = os.environ.get("PYTORCH_DEVICE") or "cpu"
YOLOV5_PATH = os.environ.get("YOLOV5_PATH")
YOLOV5_WEIGHT_PATH = os.environ.get("YOLOV5_WEIGHT_PATH")

class Operator:
    """
    Send `Right bbox and Left bbox found by YOLOv5 on given `image`
    """

    def __init__(self):
        self.COUNT_IMAGES = 1
        self.res = [[], [], []] # front, right, left
        self.point = [[], [], []]
        if YOLOV5_PATH is None:
            # With internet
            self.model = torch.hub.load(
                "ultralytics/yolov5",
                "yolov5n",
            )
        else:
            # Without internet
            self.model = torch.hub.load(
                YOLOV5_PATH,
                "custom",
                path=YOLOV5_WEIGHT_PATH,
                source="local",
            )

        # self.model.conf = 0.4 # confidence threshold
        # self.model.iou = 0.55 # NMS IoU threshold
        self.model.to(torch.device(DEVICE))
        self.model.eval()

    def on_event(
        self,
        dora_event: dict,
        send_output: Callable[[str, bytes], None],
    ) -> DoraStatus:
        if dora_event["type"] == "INPUT":
            return self.on_input(dora_event, send_output)
        return DoraStatus.CONTINUE

    def on_input(
        self,
        dora_input: dict,
        send_output: Callable[[str, bytes], None],
    ) -> DoraStatus:
        """Handle image
        Args:
            dora_input["id"](str): Id of the input declared in the yaml configuration
            dora_input["data"] (bytes): Bytes message of the input
            send_output (Callable[[str, bytes]]): Function enabling sending output back to dora.
        """
        
        if dora_input["id"] == "image":
            frames = np.frombuffer(
                dora_input["data"],
                np.uint8,
            ).reshape((3, IMAGE_HEIGHT, IMAGE_WIDTH, 4))
            frames = frames[:, :, :, :3]
            # print("count: ", self.COUNT_IMAGES)
            
            bboxes = []
            b = []
            
            for i in range(3):
                # begin = time.time()
                frame = frames[i]
                results = self.model(frame)  # includes NMS
                # print(f"yolov5_op: results: {results}")
                arrays = np.array(results.xyxy[0].cpu())[
                    :, [0, 2, 1, 3, 4, 5]
                ]  # xyxy -> xxyy
                arrays[:, 4] *= 100
                arrays = arrays.astype(np.int32)
                if arrays.shape[0] != 0:
                    if self.COUNT_IMAGES==1:
                        self.res[i].append(arrays)
                    else:
                        for each in range(len(arrays)):
                            center_label = arrays[each][-1]
                            centerP = [(arrays[each][0] + arrays[each][1])/2, (arrays[each][2] + arrays[each][3])/2]
                            if len(self.res[i]) == 0:
                                self.res[i].append(arrays)
                            else:
                                # print(np.array(self.res[i]).shape)
                                for idx in range(len(self.res[i][0])):
                                    p = self.res[i][0][idx]
                                    if p[-1] == center_label:
                                        # print(np.array(p).shape)
                                        cp = [(p[0] + p[1])/2, (p[2] + p[3])/2]
                                        # print(centerP, cp)
                                        # print(np.linalg.norm(np.array(centerP) - np.array(cp)))
                                        if np.linalg.norm(np.array(centerP) - np.array(cp)) < 10:
                                            self.res[i][0][idx][0] = (arrays[each][0] + p[0])/2
                                            self.res[i][0][idx][1] = (arrays[each][1] + p[1])/2
                                            self.res[i][0][idx][2] = (arrays[each][2] + p[2])/2
                                            self.res[i][0][idx][3] = (arrays[each][3] + p[3])/2
                                            break
                                else:
                                    # print("ap: ",np.array(self.res[i][0]).shape, np.array(arrays[each]).shape)
                                    self.res[i][0] = np.append(self.res[i][0], np.array(arrays[each]).reshape(1, -1), axis=0)
                                    # print(self.res[i][0].shape)
                # print("arrays: ", arrays.shape, np.array(self.res[i]).shape)
                # b.append(arrays)
                arrays = arrays.tobytes()
                bboxes.append(arrays)
                # print("bboxes: ", np.array(bboxes).shape)
                # end = time.time()
                # print("node2: ", end-begin)
            self.COUNT_IMAGES += 1
            # send_output("bbox", bboxes[0], dora_input["metadata"])
            # send_output("rbbox", bboxes[0], dora_input["metadata"])
            # print("left: ", b[2])
            if self.COUNT_IMAGES % 3 == 0:
                for i in range(3):
                    print(np.array(self.res[i]).shape)
                    if np.array(self.res[i]).shape[0] == 0:
                        b.append(np.zeros((0, 6)).astype(np.int32).tobytes())
                    else:
                        b.append(self.res[i][0].astype(np.int32).tobytes())
                # for i in b:
                #     print("b: ", i.shape)
                self.res = [[], [], []] # front, right, left
                self.COUNT_IMAGES = 1
            
                send_output("lbbox", np.array(bboxes).tobytes(), dora_input["metadata"])
            # send_output("lbbox", np.array(bboxes).tobytes(), dora_input["metadata"])
            return DoraStatus.CONTINUE
